import os
import gradio as gr
import agentscope
from agentscope.agents import UserAgent
from agentscope.agents.dialog_agent import DialogAgent
from agentscope.message.msg import Msg

# 初始化模型配置
model_configs = [
    {
        "model_type": "openai_chat",
        "config_name": "deepseek-r1-7b-local",
        "model_name": "deepseek-r1:7b",
        "client_args": {
            "base_url": "http://127.0.0.1:11434/v1",
        },
        "api_key": "Empty",
        "generate_args": {"temperature": 0.6},
        "stream": True,
    },
    {
        "model_type": "openai_chat",
        "config_name": "deepseek-r1:7b",
        "model_name": "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B",
        "client_args": {
            "base_url": "https://api.siliconflow.cn/v1",
        },
        "api_key": os.environ.get("SILICON_API_KEY"),
        "generate_args": {"temperature": 0.6},
        "stream": True,
    }
]

# agentscope.init(model_configs=model_configs, studio_url="http://127.0.0.1:5000/")
agentscope.init(model_configs=model_configs)


# 创建代理
agent = DialogAgent(
    name="assistant",
    # model_config_name="deepseek-r1-7b-local",
    model_config_name="deepseek-r1:7b",
    verbose=True,
    sys_prompt="You are a helpful assistant!"
)
user = UserAgent(name="User", input_hint="User Input ('exit' to quit): ")

# 定义对话函数
def chat(message, history):
    history = history or []
    msg = Msg(name="user",role="user",content=message)
    x = agent(msg)
    history.append((message, x.content))
    return x.content

# 创建 Gradio 界面
iface = gr.ChatInterface(
    fn=chat,
    title="DeepSeek-R1",
    description="与 AI 助手进行对话。输入 'exit' 结束对话。",
    examples=["请逐步推理，然后将最终答案放在\\boxed {}中。\n今有雉兔同笼，上有三十五头，下有九十四足，问雉兔各几何？", "Windows系统如何加速系统启动速度"],
)

# 启动 Gradio 应用
iface.launch()